from resnet import *
import torch.nn as nn
import torch


class ResGazeEs(nn.Module):

    def __init__(self, ):
        super(ResGazeEs, self).__init__()
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(2048, 2)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

    def forward(self, x):
        x = self.avgpool(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x


class model(nn.Module):
    def __init__(self):
        super(model, self).__init__()
        self.gaze_network = resnet50(pretrained=True)

        # self.gaze_fc = nn.Sequential(
        #     nn.Linear(2048, 2),
        # )
        self.gazeEs = ResGazeEs()

    def forward(self, x):
        feature = self.gaze_network(x)
        # feature = feature.view(feature.size(0), -1)
        # gaze = self.gaze_fc(feature)
        gaze = self.gazeEs(feature)

        return gaze


if __name__ == "__main__":
    model_test = model()
    input = torch.rand(16, 3, 224, 224)
    print(model_test(input))
